from typing import Any, Callable, Self

#"pip install PyOpenGL"
from OpenGL.GL import *
from OpenGL.constant import Constant as _C
from OpenGL.GL.ARB.shader_objects   import *
from OpenGL.GL.ARB.vertex_shader    import *
from OpenGL.GL.ARB.geometry_shader4 import *
from OpenGL.GL.ARB.fragment_shader  import *
_GLSHAD_IMPL = -1

try:
	#(another basecode file)
	import gl_texture
	_GLSHAD_HAS_TEX = True
except:
	_GLSHAD_HAS_TEX = False



def _get_processed_log( log:str|bytes ) -> str:
	if type(log) == bytes:
		return _get_processed_log( log.decode() )
	assert type(log) == str

	return "\n".join([ line for line in log.split("\n") if line ])



class ShaderBase:
	def __init__( self:Self, type:_C, src:str ):
		self.type = type

		self.source = src

		def get_handle():
			global _GLSHAD_IMPL
			if _GLSHAD_IMPL == -1:
				try:
					handle = glCreateShader( type )
					_GLSHAD_IMPL = 0
					return handle
				except: pass
				try:
					handle = glCreateShaderObjectARB( type )
					_GLSHAD_IMPL = 1
					return handle
				except: pass
				raise RuntimeError( "Shader object not available!" )
			elif _GLSHAD_IMPL == 0:
				handle = glCreateShader( type )
			else:
				handle = glCreateShaderObjectARB( type )
			return handle
		self._handle = get_handle()

		if _GLSHAD_IMPL == 0:
			glShaderSource( self._handle, [src] ) # Note must be array of strings
			glCompileShader( self._handle )

			status = glGetShaderiv( self._handle, GL_COMPILE_STATUS )
		else:
			glShaderSourceARB( self._handle, [src] ) # Note must be array of strings
			glCompileShaderARB( self._handle )

			status = glGetObjectParameterivARB( self._handle, GL_OBJECT_COMPILE_STATUS_ARB )

		if status != GL_TRUE:
			raise RuntimeError( "Shader compile failed:\n" + self.get_log_str() )

	def __del__( self:Self ):
		if _GLSHAD_IMPL == 0: glDeleteShader   ( self._handle )
		else                : glDeleteObjectARB( self._handle )

	def get_log_str( self:Self ):
		if _GLSHAD_IMPL == 0: log=glGetShaderInfoLog( self._handle )
		else                : log=glGetInfoLogARB   ( self._handle )
		return _get_processed_log( log )

class ShaderVert(ShaderBase):
	def __init__( self:Self, src:str ):
		if _GLSHAD_IMPL == 0: ShaderBase.__init__( self, GL_VERTEX_SHADER         , src )
		else                : ShaderBase.__init__( self, GL_VERTEX_SHADER_ARB     , src )

class ShaderTessCtrl(ShaderBase):
	def __init__( self:Self, src:str ):
		if _GLSHAD_IMPL == 0: ShaderBase.__init__( self, GL_TESS_CONTROL_SHADER   , src )
		else                : ShaderBase.__init__( self, GL_TESS_CONTROL_SHADER   , src ) #PyOpenGL doesn't add ARB here?

class ShaderTessEval(ShaderBase):
	def __init__( self:Self, src:str ):
		if _GLSHAD_IMPL == 0: ShaderBase.__init__( self, GL_TESS_EVALUATION_SHADER, src )
		else                : ShaderBase.__init__( self, GL_TESS_EVALUATION_SHADER, src ) #PyOpenGL doesn't add ARB here?

class ShaderGeom(ShaderBase):
	def __init__( self:Self, src:str ):
		if _GLSHAD_IMPL == 0: ShaderBase.__init__( self, GL_GEOMETRY_SHADER       , src )
		else                : ShaderBase.__init__( self, GL_GEOMETRY_SHADER_ARB   , src )

class ShaderFrag(ShaderBase):
	def __init__( self:Self, src:str ):
		if _GLSHAD_IMPL == 0: ShaderBase.__init__( self, GL_FRAGMENT_SHADER       , src )
		else                : ShaderBase.__init__( self, GL_FRAGMENT_SHADER_ARB   , src )



class Program:
	current:Self|None = None

	def __init__( self:Self, shaders:list[ShaderBase] ):
		assert len(shaders) != 0
		for shader in shaders:
			assert isinstance( shader, ShaderBase )

		self.shaders = shaders

		if _GLSHAD_IMPL == 0:
			self._handle = glCreateProgram()

			for shader in shaders:
				glAttachShader( self._handle, shader._handle )

			glLinkProgram( self._handle )
			status = glGetProgramiv( self._handle, GL_LINK_STATUS )
			if status != GL_TRUE:
				raise RuntimeError( "Program link failed:\n" + self.get_log_str() )

			glValidateProgram( self._handle )
			status = glGetProgramiv( self._handle, GL_VALIDATE_STATUS )
			if status != GL_TRUE:
				raise RuntimeError( "Program validation failed:\n" + self.get_log_str() )

		else:
			self._handle = glCreateProgramObjectARB()

			for shader in shaders:
				glAttachObjectARB( self._handle, shader._handle )

			glLinkProgramARB( self._handle )
			status = glGetObjectParameterivARB( self._handle, GL_OBJECT_LINK_STATUS_ARB )
			if status != GL_TRUE:
				raise RuntimeError( "Program link failed:\n" + self.get_log_str() )

			glValidateProgramARB( self._handle )
			status = glGetObjectParameterivARB( self._handle, GL_OBJECT_VALIDATE_STATUS_ARB )
			if status != GL_TRUE:
				raise RuntimeError( "Program validation failed:\n" + self.get_log_str() )

		self._sym_locs = {}

	def __del__( self:Self ):
		if _GLSHAD_IMPL == 0: glDeleteProgram  ( self._handle )
		else                : glDeleteObjectARB( self._handle )

	def get_log_str( self:Self ):
		if _GLSHAD_IMPL == 0: log=glGetProgramInfoLog( self._handle )
		else                : log=glGetInfoLogARB    ( self._handle )
		return _get_processed_log( log )

	def save_sources( self:Self, path:str ):
		file = open( path, "w" )
		for shader in self.shaders:
			file.write( shader.source )
			file.write( "\n" )
		file.close()

	def get_location( self:Self, symbol:str ):
		if not symbol in self._sym_locs.keys():
			if _GLSHAD_IMPL == 0: loc=glGetUniformLocation   ( self._handle, symbol )
			else                : loc=glGetUniformLocationARB( self._handle, symbol )
			if loc == -1:
				raise RuntimeError(f"Cannot get the location of symbol \"{symbol}\"!")
			self._sym_locs[ symbol ] = loc
		return self._sym_locs[ symbol ]
	
	def _pass( self:Self, symbol:str, fn:Callable,fnarb:Callable|None, *args ):
		assert Program.current == self
		loc = self.get_location( symbol )
		if _GLSHAD_IMPL == 0: fn   ( loc, *args )
		elif fnarb != None  : fnarb( loc, *args )
		else                : raise RuntimeError( "Passing this type unsupported!" )

	# Pass vectors
	def pass_1i( self:Self, symbol:str, val:int ):
		self._pass( symbol, glUniform1i,glUniform1iARB, val )
	def pass_2i( self:Self, symbol:str, vec:tuple[int,int]|list ):
		self._pass( symbol, glUniform2i,glUniform2iARB, vec[0],vec[1] )
	def pass_3i( self:Self, symbol:str, vec:tuple[int,int,int]|list ):
		self._pass( symbol, glUniform2i,glUniform2iARB, vec[0],vec[1],vec[2] )
	def pass_4i( self:Self, symbol:str, vec:tuple[int,int,int,int]|list ):
		self._pass( symbol, glUniform2i,glUniform2iARB, vec[0],vec[1],vec[2],vec[3] )
	def pass_1u( self:Self, symbol:str, val:int ):
		self.pass_1i( self, symbol, val )
	def pass_2u( self:Self, symbol:str, vec:tuple[int,int]|list ):
		self.pass_2i( self, symbol, vec )
	def pass_3u( self:Self, symbol:str, vec:tuple[int,int,int]|list ):
		self.pass_3i( self, symbol, vec )
	def pass_4u( self:Self, symbol:str, vec:tuple[int,int,int,int]|list ):
		self.pass_4i( self, symbol, vec )
	def pass_1b( self:Self, symbol:str, val:bool ):
		self.pass_1i( self, symbol, int(val) )
	def pass_2b( self:Self, symbol:str, vec:tuple[bool,bool]|list ):
		self.pass_2i( self, symbol, (int(vec[0]),int(vec[1])) )
	def pass_3b( self:Self, symbol:str, vec:tuple[bool,bool,bool]|list ):
		self.pass_3i( self, symbol, (int(vec[0]),int(vec[1]),int(vec[2])) )
	def pass_4b( self:Self, symbol:str, vec:tuple[bool,bool,bool,bool]|list ):
		self.pass_4i( self, symbol, (int(vec[0]),int(vec[1]),int(vec[2]),int(vec[3])) )
	def pass_1f( self:Self, symbol:str, val:float ):
		self._pass( symbol, glUniform1f,glUniform1fARB, val )
	def pass_2f( self:Self, symbol:str, vec:tuple[float,float]|list ):
		self._pass( symbol, glUniform2f,glUniform2fARB, vec[0],vec[1] )
	def pass_3f( self:Self, symbol:str, vec:tuple[float,float,float]|list ):
		self._pass( symbol, glUniform2f,glUniform2fARB, vec[0],vec[1],vec[2] )
	def pass_4f( self:Self, symbol:str, vec:tuple[float,float,float,float]|list ):
		self._pass( symbol, glUniform2f,glUniform2fARB, vec[0],vec[1],vec[2],vec[3] )

	# Pass matrices.  Note e.g. "3x4" means 3 rows, 4 columns, as is standard in mathematics.
	def pass_2x2f( self:Self, symbol:str, matr:Any, transpose:bool=False ):
		args = [ 1, int(transpose)[GL_FALSE,GL_TRUE] ] + Program._get_flat_matr( matr )
		self._pass( symbol, glUniformMatrix2fv  ,glUniformMatrix2fvARB, *args )
	def pass_2x3f( self:Self, symbol:str, matr:Any, transpose:bool=False ):
		args = [ 1, int(transpose)[GL_FALSE,GL_TRUE] ] + Program._get_flat_matr( matr )
		self._pass( symbol, glUniformMatrix3x2fv,None                 , *args )
	def pass_2x4f( self:Self, symbol:str, matr:Any, transpose:bool=False ):
		args = [ 1, int(transpose)[GL_FALSE,GL_TRUE] ] + Program._get_flat_matr( matr )
		self._pass( symbol, glUniformMatrix4x2fv,None                 , *args )
	def pass_3x2f( self:Self, symbol:str, matr:Any, transpose:bool=False ):
		args = [ 1, int(transpose)[GL_FALSE,GL_TRUE] ] + Program._get_flat_matr( matr )
		self._pass( symbol, glUniformMatrix2x3fv,None               , *args )
	def pass_3x3f( self:Self, symbol:str, matr:Any, transpose:bool=False ):
		args = [ 1, int(transpose)[GL_FALSE,GL_TRUE] ] + Program._get_flat_matr( matr )
		self._pass( symbol, glUniformMatrix3fv  ,glUniformMatrix3fvARB, *args )
	def pass_3x4f( self:Self, symbol:str, matr:Any, transpose:bool=False ):
		args = [ 1, int(transpose)[GL_FALSE,GL_TRUE] ] + Program._get_flat_matr( matr )
		self._pass( symbol, glUniformMatrix4x3fv,None               , *args )
	def pass_4x2f( self:Self, symbol:str, matr:Any, transpose:bool=False ):
		args = [ 1, int(transpose)[GL_FALSE,GL_TRUE] ] + Program._get_flat_matr( matr )
		self._pass( symbol, glUniformMatrix2x4fv,None               , *args )
	def pass_4x3f( self:Self, symbol:str, matr:Any, transpose:bool=False ):
		args = [ 1, int(transpose)[GL_FALSE,GL_TRUE] ] + Program._get_flat_matr( matr )
		self._pass( symbol, glUniformMatrix3x4fv,None               , *args )
	def pass_4x4f( self:Self, symbol:str, matr:Any, transpose:bool=False ):
		args = [ 1, int(transpose)[GL_FALSE,GL_TRUE] ] + Program._get_flat_matr( matr )
		self._pass( symbol, glUniformMatrix4fv  ,glUniformMatrix4fvARB, *args )

	if _GLSHAD_HAS_TEX:
		def pass_tex2D( self:Self, tex:gl_texture.Texture2D, index:int ):
			if index != 0: glActiveTexture( GL_TEXTURE0 + index )
			gl_texture.Texture2D.bind( tex )
			self.pass_1i( f"tex2D_{index}", index )
			if index != 0: glActiveTexture( GL_TEXTURE0 )

	@staticmethod
	def use( program:Self|None=None ):
		if program == Program.current: return

		handle = 0 if program==None else program._handle
		if _GLSHAD_IMPL == 0: glUseProgram         ( handle )
		else                : glUseProgramObjectARB( handle )
		Program.current = program
